from tqdm import tqdm
import os
import shelve
import numpy as np


def getDisAndLableSampleName(rootPath, labelPath):
    """
        output: [name1, name2,...]
        desc: 获得能反汇编且拥有家族标签的样本名
    """
    samples = {}
    for i in os.listdir(rootPath):
        tmp = i.split(".")
        if tmp[1] == "dir":
            samples[tmp[0]] = 1
    ans = []
    with open(labelPath, "r") as file:
        for line in file:
            tmp = line.split("\t")
            sampleName = tmp[0]
            if sampleName in samples:
                ans.append(sampleName)
    return ans


def map_api_name_to_eight_bit_vector(function_name, embeddingSize, numLim):
    """
        desc: 将api转换成一个embeddingSize位的向量,每一位使用numLim取余
    """
    function_bytes = function_name.encode('utf-8')
    eight_bit_vector = [0] * embeddingSize
    for i, byte in enumerate(function_bytes):
        index = i % embeddingSize
        eight_bit_vector[index] += byte
        eight_bit_vector[index] %= numLim
    return eight_bit_vector


def getCGData(rootPath, name):
    """
        input: 样本名
        output: 样本CG (adj,att,vType)
                    adj 邻接矩阵 b*n*n  # 似乎没有这个batch
                    att 特征矩阵 b*n*d
                    vtype 类型矩阵 b*n*3  3种类型
                    startInd int 起始函数的下标
        desc: 获得样本的CG数据
    """
    fileroot = rootPath+"//"+name
    res = {}
    res["adj"], res["att"], res["vtype"] = [], [], []
    res["isSuccess"] = True
    with shelve.open(fileroot) as file:
        cg = file["cg"]
        res["adj"] = np.array(cg)  # 这里是带边权的

        cgattr = file["cgattr"]
        try:
            functype = file["funcType"]
        except Exception as e:
            functype = file["functype"]
        try:
            funcs_id = file["func_id"]  # name-->ind
        except Exception as e:
            funcs_id = file["funcs_id"]  # name-->ind
    # print(funcs_id)
    # print(functype)
    hasStartFunc = False
    startFuncInd = -1
    id_funcs = {}
    for i in funcs_id.keys():
        id_funcs[funcs_id[i]] = i
        # if i == "StartAddress" or i == "start":
        if i == "start":
            hasStartFunc = True
            startFuncInd = funcs_id[i]
    if hasStartFunc == False:
        print("error: 不存在起始函数，样本名：{}".format(name))
        res["isSuccess"] = False
        return res

    #   动态导入函数是[],需要处理一下,
    #   同时生成vtype
    tempData = []
    for i in functype.keys():
        temp = []
        if functype[i] == "local":
            temp = [1, 0, 0]
        elif functype[i] == 'dynamic import':
            temp = [0, 1, 0]
            cgattr[i] = map_api_name_to_eight_bit_vector(id_funcs[i], 8, 20)
        else:
            temp = [0, 0, 1]
        tempData.append(temp)
    res["vType"] = np.array(tempData)
    res["att"] = np.array(cgattr)  # 应该是8维，但是第一维字符串的提取似乎出现了问题，建议只是用后七维
    res["startInd"] = startFuncInd
    return res


def normalize_adj_matrix(A):
    """
        A: 输入的邻接矩阵，二维 numpy 数组
        desc: 给出矩阵的非对称归一化矩阵
    """
    # 计算度数矩阵 D
    degree_matrix = np.diag(np.sum(A, axis=1))  # 每一行的元素之和，即度数矩阵

    # 计算 D^(-1) * A，进行行归一化
    # D_inv = np.linalg.inv(degree_matrix)  # 计算度数矩阵的逆
    D_inv = np.linalg.pinv(degree_matrix)  # 计算度数矩阵的逆

    # 归一化邻接矩阵 A' = D^(-1) * A
    A_normalized = np.dot(D_inv, A)

    return A_normalized


def k_step_neighbors(A, X, k):
    """
        计算矩阵 A 的 k 阶邻接信息，并与特征矩阵 X 相乘
        A: 邻接矩阵
        X: 特征矩阵
        k: 邻居步数
    """
    # 计算 A^k，邻接矩阵的k次幂
    A_k = np.linalg.matrix_power(A, k)
    # 将 A^k 与特征矩阵 X 相乘，得到 k 阶邻居的信息
    result = np.dot(A_k, X)
    return result


def removeNode(cgAtt, cgVType, cgAdj, delNode):
    """
        移除图结构特定的结点
    """
    EECG = {}
    size = len(cgAdj)
    EECGSize = size - len(delNode)
    adj = np.zeros((EECGSize, EECGSize))
    att = np.zeros((EECGSize, 8))
    vType = np.zeros((EECGSize, 3))

    # 新旧矩阵的下标映射
    old2new = {}
    ind = 0
    for i in range(size):
        if i not in delNode:
            old2new[i] = ind
            ind += 1

    for i in range(size):
        for j in range(size):
            if (i not in delNode and j not in delNode and cgAdj[i][j] != 0):
                adj[old2new[i]][old2new[j]] = cgAdj[i][j]
    for i in range(size):
        if i not in delNode:
            att[old2new[i]] = cgAtt[i]
            vType[old2new[i]] = cgVType[i]
    EECG["adj"] = adj
    EECG["att"] = att
    EECG["vType"] = vType
    return EECG


def CG2EECG(cgData):
    """
        input: CG
        output: EECG
        desc: 将CG转换成EECG
    """
    # 1. 实现公式4-1 && 4-2
    att = cgData["att"]
    vType = cgData["vType"]
    adj = cgData["adj"]
    startInd = cgData["startInd"]
    normalizeAdj = normalize_adj_matrix(adj)
    att_k = k_step_neighbors(normalizeAdj, att, 3)

    # 2. del all local nodes
    size = len(adj)
    delNode = {}
    for ind in range(size):
        if vType[ind][0] == 0:
            continue
        if ind == startInd:
            continue
        callers, callees = [], []
        externalInd = -1
        for i in range(size):
            if adj[i][ind] != 0:
                callers.append(i)
            if adj[ind][i] != 0:
                if externalInd == -1 and vType[i][0] != 0:
                    externalInd = i
                callees.append(i)
        # 孩子没有外部结点
        if externalInd == -1:
            for caller in callers:
                for callee in callees:
                    adj[caller][callee] = 1
        else:
            for caller in callers:
                adj[caller][externalInd] = 1
            for callee in callees:
                if callee != externalInd:
                    adj[externalInd][callee] = 1
        # remove node
        delNode[ind] = vType[ind]
        vType[ind] = [0, 0, 0]
        for j in range(size):
            adj[ind][j] = 0
            adj[j][ind] = 0
    # 3. generate EECG
    return removeNode(att_k, vType, adj, delNode)


def saveEECG(filePath, EECG):
    """
        保存EECG数据
    """
    with shelve.open(filePath) as file:
        file['adj'] = EECG["adj"]
        file['att'] = EECG["att"]
        file['vType'] = EECG["vType"]


def getEECG(filePath):
    with shelve.open(filePath) as file:
        print(file['adj'])
        print(file['att'])
        print(file['vType'])


if __name__ == "__main__":
    rootPath = r"/mnt/mydisk1/chenyongwei/sampleDatas/functionSim"
    labelPath = r"/home/cyw/projects/function_sim_project/all_data/fam_lables.txt"
    savePath = r"/mnt/mydisk1/chenyongwei/sampleDatas/EECG"

    sampleNames = getDisAndLableSampleName(rootPath, labelPath)
    for sampleName in tqdm(sampleNames):
        cgData = getCGData(rootPath, sampleName)
        if (cgData["isSuccess"]):
            eecgData = CG2EECG(cgData)
            saveEECG(savePath+"//"+sampleName, eecgData)
